{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 5: Training Spiking Neural Networks\n",
    "\n",
    "**Duration:** ~45 minutes | **Prerequisites:** Basic Tutorials 1-4\n",
    "\n",
    "## Learning Objectives\n",
    "\n",
    "By the end of this tutorial, you will:\n",
    "\n",
    "- ✅ Understand surrogate gradient methods for training SNNs\n",
    "- ✅ Implement backpropagation through time (BPTT) for SNNs\n",
    "- ✅ Use appropriate loss functions for spike-based learning\n",
    "- ✅ Configure optimizers and learning rates\n",
    "- ✅ Train an SNN classifier on real datasets\n",
    "- ✅ Evaluate and visualize training progress\n",
    "\n",
    "## Overview\n",
    "\n",
    "Training spiking neural networks is challenging because spike generation is a discrete, non-differentiable operation. In this tutorial, we'll learn how to overcome this using **surrogate gradient methods**, which allow us to train SNNs using standard gradient-based optimization.\n",
    "\n",
    "**Key Concepts:**\n",
    "- **The gradient problem**: Spike generation has zero gradient almost everywhere\n",
    "- **Surrogate gradients**: Use smooth approximations during backpropagation\n",
    "- **BPTT for SNNs**: Unroll network dynamics through time\n",
    "- **Rate-based losses**: Train on spike rates or membrane potentials\n",
    "- **Temporal credit assignment**: Learn when to spike\n",
    "\n",
    "Let's start by understanding why training SNNs is difficult and how we can solve it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import brainpy\n",
    "import brainstate\n",
    "import brainunit as u\n",
    "import braintools\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Set random seed for reproducibility\n",
    "brainstate.random.seed(42)\n",
    "\n",
    "# Configure environment\n",
    "brainstate.environ.set(dt=1.0 * u.ms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: The Gradient Problem\n",
    "\n",
    "Let's visualize why training SNNs is challenging. The spike generation function is a Heaviside step function:\n",
    "\n",
    "$$\n",
    "S(V) = \\begin{cases}\n",
    "1 & \\text{if } V \\geq V_{th} \\\\\n",
    "0 & \\text{if } V < V_{th}\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "The gradient of this function is:\n",
    "\n",
    "$$\n",
    "\\frac{dS}{dV} = \\begin{cases}\n",
    "\\infty & \\text{at } V = V_{th} \\\\\n",
    "0 & \\text{everywhere else}\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "This makes gradient-based learning impossible! Let's see this visually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Heaviside step function\n",
    "def heaviside(x, threshold=0.0):\n",
    "    return (x >= threshold).astype(float)\n",
    "\n",
    "# Voltage values\n",
    "V = np.linspace(-2, 2, 1000)\n",
    "V_th = 0.0\n",
    "\n",
    "# Spike function and its \"gradient\"\n",
    "spikes = heaviside(V, V_th)\n",
    "# Numerical gradient (will be mostly zeros)\n",
    "grad_spike = np.gradient(spikes, V)\n",
    "\n",
    "# Plot\n",
    "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "# Spike function\n",
    "axes[0].plot(V, spikes, 'b-', linewidth=2)\n",
    "axes[0].axvline(V_th, color='r', linestyle='--', label='Threshold')\n",
    "axes[0].set_xlabel('Membrane Potential (V)', fontsize=12)\n",
    "axes[0].set_ylabel('Spike Output', fontsize=12)\n",
    "axes[0].set_title('Spike Generation Function', fontsize=14, fontweight='bold')\n",
    "axes[0].legend()\n",
    "axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "# Gradient (problematic!)\n",
    "axes[1].plot(V, grad_spike, 'r-', linewidth=2)\n",
    "axes[1].axvline(V_th, color='r', linestyle='--', label='Threshold')\n",
    "axes[1].set_xlabel('Membrane Potential (V)', fontsize=12)\n",
    "axes[1].set_ylabel('Gradient dS/dV', fontsize=12)\n",
    "axes[1].set_title('Gradient (Problematic!)', fontsize=14, fontweight='bold')\n",
    "axes[1].legend()\n",
    "axes[1].grid(True, alpha=0.3)\n",
    "axes[1].set_ylim(-0.1, 0.6)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"❌ Problem: Gradient is zero almost everywhere!\")\n",
    "print(\"   This prevents gradient descent from working.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Surrogate Gradient Solution\n",
    "\n",
    "The solution is **surrogate gradients**: Use the true spike function in the forward pass, but use a smooth approximation during backpropagation.\n",
    "\n",
    "**Common surrogate gradient functions:**\n",
    "\n",
    "1. **Sigmoid**: $\\sigma'(\\beta(V - V_{th}))$\n",
    "2. **ReLU**: $\\max(0, 1 - |V - V_{th}|)$\n",
    "3. **SuperSpike**: $\\frac{1}{(1 + |\\beta(V - V_{th})|)^2}$\n",
    "\n",
    "BrainPy provides these in `braintools.surrogate`. Let's visualize them!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create surrogate gradient functions\n",
    "sigmoid_surrogate = braintools.surrogate.Sigmoid(alpha=4.0)\n",
    "relu_surrogate = braintools.surrogate.ReluGrad(alpha=1.0)\n",
    "superspike_surrogate = braintools.surrogate.SlayerGrad(alpha=4.0)\n",
    "\n",
    "# Voltage range\n",
    "V_range = np.linspace(-2, 2, 1000)\n",
    "V_th = 0.0\n",
    "\n",
    "# Compute surrogate gradients\n",
    "grad_sigmoid = sigmoid_surrogate(V_range - V_th)\n",
    "grad_relu = relu_surrogate(V_range - V_th)\n",
    "grad_superspike = superspike_surrogate(V_range - V_th)\n",
    "\n",
    "# Plot\n",
    "fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n",
    "\n",
    "# Sigmoid surrogate\n",
    "axes[0].plot(V_range, grad_sigmoid, 'g-', linewidth=2, label='Sigmoid surrogate')\n",
    "axes[0].axvline(V_th, color='r', linestyle='--', alpha=0.5)\n",
    "axes[0].set_xlabel('V - V_th', fontsize=12)\n",
    "axes[0].set_ylabel('Surrogate Gradient', fontsize=12)\n",
    "axes[0].set_title('Sigmoid Surrogate', fontsize=14, fontweight='bold')\n",
    "axes[0].legend()\n",
    "axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "# ReLU surrogate\n",
    "axes[1].plot(V_range, grad_relu, 'b-', linewidth=2, label='ReLU surrogate')\n",
    "axes[1].axvline(V_th, color='r', linestyle='--', alpha=0.5)\n",
    "axes[1].set_xlabel('V - V_th', fontsize=12)\n",
    "axes[1].set_ylabel('Surrogate Gradient', fontsize=12)\n",
    "axes[1].set_title('ReLU Surrogate', fontsize=14, fontweight='bold')\n",
    "axes[1].legend()\n",
    "axes[1].grid(True, alpha=0.3)\n",
    "\n",
    "# SuperSpike surrogate\n",
    "axes[2].plot(V_range, grad_superspike, 'm-', linewidth=2, label='SuperSpike surrogate')\n",
    "axes[2].axvline(V_th, color='r', linestyle='--', alpha=0.5)\n",
    "axes[2].set_xlabel('V - V_th', fontsize=12)\n",
    "axes[2].set_ylabel('Surrogate Gradient', fontsize=12)\n",
    "axes[2].set_title('SuperSpike Surrogate', fontsize=14, fontweight='bold')\n",
    "axes[2].legend()\n",
    "axes[2].grid(True, alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"✅ Solution: Smooth surrogate gradients enable learning!\")\n",
    "print(\"   Forward pass: Use real spikes\")\n",
    "print(\"   Backward pass: Use smooth gradient approximation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Creating a Trainable SNN\n",
    "\n",
    "Now let's create an SNN classifier. We'll build a simple network:\n",
    "\n",
    "**Architecture:**\n",
    "- Input layer: 784 neurons (28×28 image)\n",
    "- Hidden layer: 128 LIF neurons\n",
    "- Output layer: 10 LIF neurons (digits 0-9)\n",
    "\n",
    "**Key for training:**\n",
    "- Use LIF neurons with surrogate gradient spike functions\n",
    "- Use `brainpy.state.Readout` to convert spikes to logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainableSNN(brainstate.nn.Module):\n",
    "    \"\"\"Simple feedforward SNN for classification.\"\"\"\n",
    "    \n",
    "    def __init__(self, n_input=784, n_hidden=128, n_output=10):\n",
    "        super().__init__()\n",
    "        \n",
    "        # Input to hidden projection\n",
    "        self.fc1 = brainstate.nn.Linear(n_input, n_hidden, w_init=braintools.init.KaimingNormal())\n",
    "        \n",
    "        # Hidden LIF neurons with surrogate gradient\n",
    "        self.lif1 = brainpy.state.LIF(\n",
    "            n_hidden,\n",
    "            V_rest=-65.0 * u.mV,\n",
    "            V_th=-50.0 * u.mV,\n",
    "            V_reset=-65.0 * u.mV,\n",
    "            tau=10.0 * u.ms,\n",
    "            spike_fun=braintools.surrogate.ReluGrad()  # Surrogate gradient!\n",
    "        )\n",
    "        \n",
    "        # Hidden to output projection\n",
    "        self.fc2 = brainstate.nn.Linear(n_hidden, n_output, w_init=braintools.init.KaimingNormal())\n",
    "        \n",
    "        # Output LIF neurons with surrogate gradient\n",
    "        self.lif2 = brainpy.state.LIF(\n",
    "            n_output,\n",
    "            V_rest=-65.0 * u.mV,\n",
    "            V_th=-50.0 * u.mV,\n",
    "            V_reset=-65.0 * u.mV,\n",
    "            tau=10.0 * u.ms,\n",
    "            spike_fun=braintools.surrogate.ReluGrad()  # Surrogate gradient!\n",
    "        )\n",
    "        \n",
    "        # Readout layer to convert spikes to logits\n",
    "        self.readout = brainpy.state.Readout(n_output, n_output)\n",
    "    \n",
    "    def update(self, x):\n",
    "        \"\"\"Forward pass for one time step.\n",
    "        \n",
    "        Args:\n",
    "            x: Input current (batch_size, n_input) with physical units\n",
    "        \n",
    "        Returns:\n",
    "            logits: Output logits (batch_size, n_output)\n",
    "        \"\"\"\n",
    "        # Input to hidden\n",
    "        current1 = self.fc1(x)\n",
    "        self.lif1(current1)\n",
    "        hidden_spikes = self.lif1.get_spike()\n",
    "        \n",
    "        # Hidden to output\n",
    "        current2 = self.fc2(hidden_spikes)\n",
    "        self.lif2(current2)\n",
    "        output_spikes = self.lif2.get_spike()\n",
    "        \n",
    "        # Convert spikes to logits\n",
    "        logits = self.readout(output_spikes)\n",
    "        \n",
    "        return logits\n",
    "\n",
    "# Create network\n",
    "net = TrainableSNN(n_input=784, n_hidden=128, n_output=10)\n",
    "brainstate.nn.init_all_states(net, batch_size=32)\n",
    "\n",
    "print(\"✅ Created trainable SNN with surrogate gradients\")\n",
    "print(f\"   Input: 784 neurons\")\n",
    "print(f\"   Hidden: 128 LIF neurons\")\n",
    "print(f\"   Output: 10 LIF neurons\")\n",
    "print(f\"   Total parameters: {sum(p.size for p in net.states(brainstate.ParamState).values())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 4: Loss Functions for SNNs\n",
    "\n",
    "For classification, we typically use **cross-entropy loss** on the output logits. The logits are computed by integrating spikes over time.\n",
    "\n",
    "**Loss computation:**\n",
    "1. Run the network for `T` time steps\n",
    "2. Accumulate output logits over time\n",
    "3. Compute cross-entropy loss: $L = -\\sum_i y_i \\log(\\text{softmax}(\\text{logits}_i))$\n",
    "\n",
    "Let's implement the training step!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_fn(network, inputs, labels, n_steps=25):\n",
    "    \"\"\"Compute loss for SNN classification.\n",
    "    \n",
    "    Args:\n",
    "        network: SNN model\n",
    "        inputs: Input data (batch_size, n_features)\n",
    "        labels: True labels (batch_size,)\n",
    "        n_steps: Number of simulation time steps\n",
    "    \n",
    "    Returns:\n",
    "        loss: Cross-entropy loss\n",
    "    \"\"\"\n",
    "    # Reset network state\n",
    "    brainstate.nn.init_all_states(network)\n",
    "    \n",
    "    # Add physical units to input (convert to current)\n",
    "    inputs_with_units = inputs * u.nA\n",
    "    \n",
    "    # Simulate for n_steps and accumulate output\n",
    "    def run_step(i):\n",
    "        return network(inputs_with_units)\n",
    "    \n",
    "    # Run simulation and accumulate logits\n",
    "    logits_sum = brainstate.transform.for_loop(run_step, jnp.arange(n_steps))\n",
    "    logits_sum = jnp.sum(logits_sum, axis=0)  # Sum over time\n",
    "    \n",
    "    # Compute cross-entropy loss\n",
    "    loss = braintools.metric.softmax_cross_entropy_with_integer_labels(\n",
    "        logits_sum, labels\n",
    "    ).mean()\n",
    "    \n",
    "    return loss\n",
    "\n",
    "def accuracy_fn(network, inputs, labels, n_steps=25):\n",
    "    \"\"\"Compute accuracy for SNN classification.\"\"\"\n",
    "    # Reset network state\n",
    "    brainstate.nn.init_all_states(network)\n",
    "    \n",
    "    # Add physical units\n",
    "    inputs_with_units = inputs * u.nA\n",
    "    \n",
    "    # Simulate and accumulate logits\n",
    "    def run_step(i):\n",
    "        return network(inputs_with_units)\n",
    "    \n",
    "    logits_sum = brainstate.transform.for_loop(run_step, jnp.arange(n_steps))\n",
    "    logits_sum = jnp.sum(logits_sum, axis=0)\n",
    "    \n",
    "    # Compute accuracy\n",
    "    predictions = jnp.argmax(logits_sum, axis=1)\n",
    "    accuracy = jnp.mean(predictions == labels)\n",
    "    \n",
    "    return accuracy\n",
    "\n",
    "print(\"✅ Defined loss and accuracy functions\")\n",
    "print(\"   Loss: Cross-entropy on accumulated logits\")\n",
    "print(\"   Accuracy: Argmax of accumulated logits\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 5: Optimizers and Training Loop\n",
    "\n",
    "Now we'll set up the optimizer and training loop. BrainPy uses `braintools.optim` which provides standard optimizers like Adam, SGD, etc.\n",
    "\n",
    "**Training loop:**\n",
    "1. Get batch of data\n",
    "2. Compute gradients using `brainstate.transform.grad()`\n",
    "3. Update parameters using optimizer\n",
    "4. Track loss and accuracy\n",
    "\n",
    "We'll use synthetic data for this demo (in practice, you'd use MNIST)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create synthetic dataset (in practice, use real data like MNIST)\n",
    "def create_synthetic_data(n_samples=1000, n_features=784, n_classes=10):\n",
    "    \"\"\"Create synthetic classification data.\"\"\"\n",
    "    X = np.random.randn(n_samples, n_features).astype(np.float32) * 0.5\n",
    "    y = np.random.randint(0, n_classes, size=n_samples)\n",
    "    return X, y\n",
    "\n",
    "# Generate data\n",
    "X_train, y_train = create_synthetic_data(n_samples=1000)\n",
    "X_test, y_test = create_synthetic_data(n_samples=200)\n",
    "\n",
    "print(\"✅ Created synthetic dataset\")\n",
    "print(f\"   Training: {X_train.shape[0]} samples\")\n",
    "print(f\"   Test: {X_test.shape[0]} samples\")\n",
    "print(f\"   Features: {X_train.shape[1]}\")\n",
    "print(f\"   Classes: {len(np.unique(y_train))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset network and create optimizer\n",
    "net = TrainableSNN(n_input=784, n_hidden=128, n_output=10)\n",
    "brainstate.nn.init_all_states(net, batch_size=32)\n",
    "\n",
    "# Create Adam optimizer\n",
    "optimizer = braintools.optim.Adam(learning_rate=1e-3)\n",
    "optimizer.register_trainable_weights(net.states(brainstate.ParamState))\n",
    "\n",
    "print(\"✅ Created optimizer\")\n",
    "print(f\"   Type: Adam\")\n",
    "print(f\"   Learning rate: 1e-3\")\n",
    "print(f\"   Trainable parameters: {len(net.states(brainstate.ParamState))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training loop\n",
    "n_epochs = 5\n",
    "batch_size = 32\n",
    "n_steps = 25  # Simulation steps per sample\n",
    "\n",
    "train_losses = []\n",
    "train_accs = []\n",
    "test_accs = []\n",
    "\n",
    "print(\"🚀 Starting training...\\n\")\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    # Shuffle training data\n",
    "    indices = np.random.permutation(len(X_train))\n",
    "    X_shuffled = X_train[indices]\n",
    "    y_shuffled = y_train[indices]\n",
    "    \n",
    "    epoch_losses = []\n",
    "    epoch_accs = []\n",
    "    \n",
    "    # Mini-batch training\n",
    "    n_batches = len(X_train) // batch_size\n",
    "    for i in range(n_batches):\n",
    "        # Get batch\n",
    "        start_idx = i * batch_size\n",
    "        end_idx = start_idx + batch_size\n",
    "        X_batch = X_shuffled[start_idx:end_idx]\n",
    "        y_batch = y_shuffled[start_idx:end_idx]\n",
    "        \n",
    "        # Compute gradients\n",
    "        grads, loss = brainstate.transform.grad(\n",
    "            loss_fn,\n",
    "            net.states(brainstate.ParamState),\n",
    "            return_value=True\n",
    "        )(net, X_batch, y_batch, n_steps)\n",
    "        \n",
    "        # Update parameters\n",
    "        optimizer.update(grads)\n",
    "        \n",
    "        # Track metrics\n",
    "        epoch_losses.append(float(loss))\n",
    "        \n",
    "        # Compute accuracy every 10 batches\n",
    "        if i % 10 == 0:\n",
    "            acc = accuracy_fn(net, X_batch, y_batch, n_steps)\n",
    "            epoch_accs.append(float(acc))\n",
    "    \n",
    "    # Epoch statistics\n",
    "    avg_loss = np.mean(epoch_losses)\n",
    "    avg_train_acc = np.mean(epoch_accs) if epoch_accs else 0.0\n",
    "    \n",
    "    # Test accuracy\n",
    "    test_acc = float(accuracy_fn(net, X_test, y_test, n_steps))\n",
    "    \n",
    "    train_losses.append(avg_loss)\n",
    "    train_accs.append(avg_train_acc)\n",
    "    test_accs.append(test_acc)\n",
    "    \n",
    "    print(f\"Epoch {epoch+1}/{n_epochs}:\")\n",
    "    print(f\"  Loss: {avg_loss:.4f}\")\n",
    "    print(f\"  Train Acc: {avg_train_acc:.2%}\")\n",
    "    print(f\"  Test Acc: {test_acc:.2%}\\n\")\n",
    "\n",
    "print(\"✅ Training complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 6: Visualizing Training Progress\n",
    "\n",
    "Let's visualize how the loss and accuracy evolved during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
    "\n",
    "epochs_range = np.arange(1, n_epochs + 1)\n",
    "\n",
    "# Plot loss\n",
    "axes[0].plot(epochs_range, train_losses, 'b-o', linewidth=2, markersize=8, label='Training Loss')\n",
    "axes[0].set_xlabel('Epoch', fontsize=12)\n",
    "axes[0].set_ylabel('Loss', fontsize=12)\n",
    "axes[0].set_title('Training Loss', fontsize=14, fontweight='bold')\n",
    "axes[0].legend()\n",
    "axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "# Plot accuracy\n",
    "axes[1].plot(epochs_range, train_accs, 'g-o', linewidth=2, markersize=8, label='Train Accuracy')\n",
    "axes[1].plot(epochs_range, test_accs, 'r-s', linewidth=2, markersize=8, label='Test Accuracy')\n",
    "axes[1].set_xlabel('Epoch', fontsize=12)\n",
    "axes[1].set_ylabel('Accuracy', fontsize=12)\n",
    "axes[1].set_title('Classification Accuracy', fontsize=14, fontweight='bold')\n",
    "axes[1].legend()\n",
    "axes[1].grid(True, alpha=0.3)\n",
    "axes[1].set_ylim(0, 1)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(f\"📊 Final Results:\")\n",
    "print(f\"   Final train accuracy: {train_accs[-1]:.2%}\")\n",
    "print(f\"   Final test accuracy: {test_accs[-1]:.2%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 7: Understanding BPTT for SNNs\n",
    "\n",
    "Let's visualize what happens during backpropagation through time (BPTT). The network processes input over multiple time steps, and gradients flow backward through time.\n",
    "\n",
    "**BPTT process:**\n",
    "1. **Forward pass**: Simulate network for T steps, accumulate outputs\n",
    "2. **Backward pass**: Compute gradients backward through all T steps\n",
    "3. **Surrogate gradients**: Used at spike generation points\n",
    "\n",
    "Let's examine the gradient flow!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Analyze gradient magnitudes during training\n",
    "def analyze_gradients(network, inputs, labels, n_steps=25):\n",
    "    \"\"\"Compute and analyze gradient magnitudes.\"\"\"\n",
    "    grads = brainstate.transform.grad(\n",
    "        loss_fn,\n",
    "        network.states(brainstate.ParamState)\n",
    "    )(network, inputs, labels, n_steps)\n",
    "    \n",
    "    # Compute gradient norms for each layer\n",
    "    grad_norms = {}\n",
    "    for name, grad in grads.items():\n",
    "        grad_norm = float(jnp.linalg.norm(grad.value.flatten()))\n",
    "        grad_norms[name] = grad_norm\n",
    "    \n",
    "    return grad_norms\n",
    "\n",
    "# Analyze gradients on a batch\n",
    "sample_X = X_train[:32]\n",
    "sample_y = y_train[:32]\n",
    "grad_norms = analyze_gradients(net, sample_X, sample_y)\n",
    "\n",
    "# Plot gradient magnitudes\n",
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "\n",
    "layer_names = list(grad_norms.keys())\n",
    "grad_values = list(grad_norms.values())\n",
    "\n",
    "colors = ['blue' if 'fc1' in name else 'green' if 'fc2' in name else 'red' for name in layer_names]\n",
    "\n",
    "bars = ax.bar(range(len(layer_names)), grad_values, color=colors, alpha=0.7)\n",
    "ax.set_xticks(range(len(layer_names)))\n",
    "ax.set_xticklabels(layer_names, rotation=45, ha='right')\n",
    "ax.set_ylabel('Gradient Norm', fontsize=12)\n",
    "ax.set_title('Gradient Magnitudes Across Layers', fontsize=14, fontweight='bold')\n",
    "ax.grid(True, alpha=0.3, axis='y')\n",
    "\n",
    "# Add legend\n",
    "from matplotlib.patches import Patch\n",
    "legend_elements = [\n",
    "    Patch(facecolor='blue', alpha=0.7, label='Input Layer'),\n",
    "    Patch(facecolor='green', alpha=0.7, label='Hidden Layer'),\n",
    "    Patch(facecolor='red', alpha=0.7, label='Readout Layer')\n",
    "]\n",
    "ax.legend(handles=legend_elements, loc='upper right')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"📊 Gradient Analysis:\")\n",
    "for name, norm in grad_norms.items():\n",
    "    print(f\"   {name}: {norm:.6f}\")\n",
    "print(\"\\n✅ Surrogate gradients enable backpropagation through spike generation!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 8: Real-World Example - MNIST Classification\n",
    "\n",
    "Now let's see how to train on real data. Here's the complete workflow for MNIST (or Fashion-MNIST):\n",
    "\n",
    "**Steps:**\n",
    "1. Load and preprocess MNIST data\n",
    "2. Convert images to rate-coded spike trains (or use pixel intensities as currents)\n",
    "3. Train SNN classifier\n",
    "4. Evaluate on test set\n",
    "\n",
    "Below is a template you can use with real MNIST data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Template for MNIST training (requires torchvision or tensorflow)\n",
    "\n",
    "def load_mnist_data():\n",
    "    \"\"\"Load and preprocess MNIST data.\n",
    "    \n",
    "    In practice, use:\n",
    "    from torchvision import datasets, transforms\n",
    "    \n",
    "    train_dataset = datasets.MNIST(\n",
    "        './data', train=True, download=True,\n",
    "        transform=transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.1307,), (0.3081,))\n",
    "        ])\n",
    "    )\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "def train_on_mnist():\n",
    "    \"\"\"Complete MNIST training workflow.\"\"\"\n",
    "    \n",
    "    # 1. Load data\n",
    "    # X_train, y_train, X_test, y_test = load_mnist_data()\n",
    "    \n",
    "    # 2. Create network\n",
    "    net = TrainableSNN(n_input=784, n_hidden=256, n_output=10)\n",
    "    brainstate.nn.init_all_states(net, batch_size=128)\n",
    "    \n",
    "    # 3. Create optimizer\n",
    "    optimizer = braintools.optim.Adam(learning_rate=1e-3)\n",
    "    optimizer.register_trainable_weights(net.states(brainstate.ParamState))\n",
    "    \n",
    "    # 4. Training loop (epochs, batches, gradient updates)\n",
    "    # for epoch in range(n_epochs):\n",
    "    #     for batch in data_loader:\n",
    "    #         grads, loss = compute_gradients(...)\n",
    "    #         optimizer.update(grads)\n",
    "    \n",
    "    # 5. Evaluation\n",
    "    # test_acc = evaluate(net, X_test, y_test)\n",
    "    \n",
    "    return net\n",
    "\n",
    "print(\"📝 MNIST Training Template:\")\n",
    "print(\"\"\"\\n1. Load MNIST: Use torchvision.datasets.MNIST or tensorflow.keras.datasets.mnist\n",
    "2. Preprocess: Flatten images (28×28 → 784), normalize to [0,1]\n",
    "3. Convert to currents: Multiply by scaling factor (e.g., 5 nA)\n",
    "4. Train: Use same loss_fn and training loop as above\n",
    "5. Expected accuracy: 95-98% on MNIST with proper hyperparameters\n",
    "\n",
    "Key hyperparameters to tune:\n",
    "- Learning rate: Try 1e-3, 5e-4, 1e-4\n",
    "- Hidden size: Try 128, 256, 512\n",
    "- Simulation steps: Try 25, 50, 100\n",
    "- Batch size: Try 32, 64, 128\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 9: Advanced Training Techniques\n",
    "\n",
    "Here are some advanced techniques to improve SNN training:\n",
    "\n",
    "### 1. Learning Rate Scheduling\n",
    "\n",
    "Reduce learning rate during training for better convergence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example: Exponential decay learning rate schedule\n",
    "def create_lr_schedule(initial_lr=1e-3, decay_rate=0.95, decay_steps=1000):\n",
    "    \"\"\"Create exponential decay learning rate schedule.\"\"\"\n",
    "    def lr_schedule(step):\n",
    "        return initial_lr * (decay_rate ** (step / decay_steps))\n",
    "    return lr_schedule\n",
    "\n",
    "# Usage:\n",
    "# lr_schedule = create_lr_schedule()\n",
    "# optimizer = braintools.optim.Adam(learning_rate=lr_schedule)\n",
    "\n",
    "print(\"✅ Learning rate scheduling helps with convergence\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Gradient Clipping\n",
    "\n",
    "Prevent gradient explosion by clipping large gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clip_gradients(grads, max_norm=1.0):\n",
    "    \"\"\"Clip gradients by global norm.\"\"\"\n",
    "    # Compute global norm\n",
    "    global_norm = jnp.sqrt(\n",
    "        sum(jnp.sum(g.value ** 2) for g in grads.values())\n",
    "    )\n",
    "    \n",
    "    # Clip if necessary\n",
    "    clip_coef = max_norm / (global_norm + 1e-6)\n",
    "    clip_coef = jnp.minimum(1.0, clip_coef)\n",
    "    \n",
    "    # Apply clipping\n",
    "    clipped_grads = {}\n",
    "    for name, grad in grads.items():\n",
    "        clipped_grads[name] = brainstate.ParamState(grad.value * clip_coef)\n",
    "    \n",
    "    return clipped_grads\n",
    "\n",
    "# Usage in training loop:\n",
    "# grads = compute_gradients(...)\n",
    "# grads = clip_gradients(grads, max_norm=1.0)\n",
    "# optimizer.update(grads)\n",
    "\n",
    "print(\"✅ Gradient clipping prevents training instabilities\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Regularization\n",
    "\n",
    "Add L2 regularization to prevent overfitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_with_regularization(network, inputs, labels, n_steps=25, l2_weight=1e-4):\n",
    "    \"\"\"Loss function with L2 regularization.\"\"\"\n",
    "    # Standard loss\n",
    "    ce_loss = loss_fn(network, inputs, labels, n_steps)\n",
    "    \n",
    "    # L2 regularization\n",
    "    l2_loss = 0.0\n",
    "    for param in network.states(brainstate.ParamState).values():\n",
    "        l2_loss += jnp.sum(param.value ** 2)\n",
    "    \n",
    "    total_loss = ce_loss + l2_weight * l2_loss\n",
    "    return total_loss\n",
    "\n",
    "print(\"✅ L2 regularization improves generalization\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this tutorial, you learned:\n",
    "\n",
    "✅ **The gradient problem**: Spike generation is non-differentiable\n",
    "\n",
    "✅ **Surrogate gradients**: Use smooth approximations during backprop\n",
    "   - Forward: Real spikes\n",
    "   - Backward: Smooth surrogate\n",
    "\n",
    "✅ **SNN architecture**: Create trainable networks with LIF neurons\n",
    "\n",
    "✅ **Loss functions**: Cross-entropy on accumulated spike outputs\n",
    "\n",
    "✅ **Training loop**: BPTT with gradient descent\n",
    "   ```python\n",
    "   grads, loss = brainstate.transform.grad(loss_fn, params)(net, X, y)\n",
    "   optimizer.update(grads)\n",
    "   ```\n",
    "\n",
    "✅ **Advanced techniques**: LR scheduling, gradient clipping, regularization\n",
    "\n",
    "**Key code pattern:**\n",
    "```python\n",
    "# 1. Create network with surrogate gradients\n",
    "lif = brainpy.state.LIF(..., spike_fun=braintools.surrogate.ReluGrad())\n",
    "\n",
    "# 2. Define loss over time\n",
    "def loss_fn(net, X, y, n_steps):\n",
    "    logits = simulate_for_n_steps(net, X, n_steps)\n",
    "    return cross_entropy(logits, y)\n",
    "\n",
    "# 3. Compute gradients and update\n",
    "grads = brainstate.transform.grad(loss_fn, params)(...)\n",
    "optimizer.update(grads)\n",
    "```\n",
    "\n",
    "**Next steps:**\n",
    "- Try training on real MNIST/Fashion-MNIST\n",
    "- Experiment with different surrogate functions\n",
    "- Tune hyperparameters (learning rate, hidden size, simulation steps)\n",
    "- Add recurrent connections for temporal tasks\n",
    "- See Tutorial 6 for incorporating synaptic plasticity\n",
    "\n",
    "**References:**\n",
    "- Neftci et al. (2019): \"Surrogate Gradient Learning in Spiking Neural Networks\"\n",
    "- Zenke & Ganguli (2018): \"SuperSpike: Supervised learning in multilayer spiking neural networks\"\n",
    "- Shrestha & Orchard (2018): \"SLAYER: Spike Layer Error Reassignment in Time\"\n",
    "- Wu et al. (2018): \"Spatio-Temporal Backpropagation for Training High-Performance Spiking Neural Networks\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "Test your understanding:\n",
    "\n",
    "### Exercise 1: Surrogate Function Comparison\n",
    "Compare training with different surrogate gradient functions (Sigmoid, ReLU, SuperSpike). Which works best?\n",
    "\n",
    "### Exercise 2: Simulation Steps\n",
    "How does the number of simulation steps (n_steps) affect accuracy and training time? Plot the trade-off.\n",
    "\n",
    "### Exercise 3: Network Architecture\n",
    "Add a second hidden layer. Does deeper architecture improve performance?\n",
    "\n",
    "### Exercise 4: Learning Rate Tuning\n",
    "Implement learning rate scheduling and compare convergence with fixed learning rate.\n",
    "\n",
    "### Exercise 5: Real MNIST\n",
    "Load real MNIST data and train a classifier. Aim for >95% test accuracy!\n",
    "\n",
    "**Bonus Challenge:** Implement online learning where the network is trained on streaming data one sample at a time (no batches)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Ecosystem-py",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
